import argparse

from nl2sql.ai_query import AIQueryManager
from nl2sql.ollama_query import OllamaQueryManager
import json
import traceback
from datetime import datetime, date
from decimal import Decimal
import os
from dotenv import load_dotenv
from nl2sql import NL2SQL

# Load environment variables from .env file
load_dotenv()

class DateTimeEncoder(json.JSONEncoder):
    """
    自定义JSON编码器，支持datetime、date和Decimal类型
    """
    def default(self, obj):
        if isinstance(obj, (datetime, date)):
            return obj.isoformat()
        if isinstance(obj, Decimal):
            return str(obj)
        return super().default(obj)

def save_result_to_file(question: str, result: dict, output_dir: str):
    """
    将查询结果保存到文件
    """
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)
    
    # 生成文件名（使用时间戳避免重名）
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"query_result_{timestamp}.txt"
    filepath = os.path.join(output_dir, filename)
    
    # 准备要保存的内容
    content = [
        f"问题: {question}",
        "-" * 50,
    ]
    
    if result['success']:
        content.extend([
            "生成的SQL:",
            result.get('sql', '无SQL信息'),
            "\n查询结果:",
            json.dumps(result['results'], ensure_ascii=False, indent=2, cls=DateTimeEncoder)
        ])
    else:
        content.extend([
            f"查询失败: {result['error']}",
            "\n错误详情:",
            result.get('traceback', '无详细信息')
        ])
    
    content.append("-" * 50)
    
    # 写入文件
    with open(filepath, 'w', encoding='utf-8') as f:
        f.write('\n'.join(content))
    
    return filepath

def print_result(question: str, result: dict, output_dir: str = None):
    """
    打印查询结果并保存到文件（如果指定了输出目录）
    """
    print("-" * 50)
    
    if result['success']:
        print("\n查询结果:")
        print(json.dumps(result['results'], ensure_ascii=False, indent=2, cls=DateTimeEncoder))
    else:
        print(f"查询失败: {result['error']}")
        if 'traceback' in result:
            print("\n错误详情:")
            print(result['traceback'])
    
    print("-" * 50)
    
    # 如果指定了输出目录，保存结果到文件
    if output_dir:
        filepath = save_result_to_file(question, result, output_dir)
        print(f"\n结果已保存到文件: {filepath}")

def save_sql_to_file(question: str, sql: str, output_dir: str):
    """
    将生成的SQL语句保存到文件
    """
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)
    
    # 生成文件名（使用时间戳避免重名）
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"sql_result_{timestamp}.txt"
    filepath = os.path.join(output_dir, filename)
    
    # 准备要保存的内容
    content = [
        f"问题: {question}",
        "-" * 50,
        "生成的SQL:",
        sql,
        "-" * 50
    ]
    
    # 写入文件
    with open(filepath, 'w', encoding='utf-8') as f:
        f.write('\n'.join(content))
    
    return filepath

def main():
    try:
        # 解析命令行参数
        parser = argparse.ArgumentParser(description='自然语言转SQL查询工具')
        parser.add_argument('--model', type=str, choices=['deepseek', 'ollama'],
                            help='选择使用的模型服务: deepseek 或 ollama (默认)')
        parser.add_argument('--interactive', action='store_true', 
                            help='交互式选择模型服务')
        args = parser.parse_args()
        
        # 如果指定了交互式模式或未指定使用的模型服务，让用户选择模型
        if args.interactive or args.model is None:
            print("请选择要使用的模型服务:")
            print("1. DeepSeek API")
            print("2. 本地OLLAMA服务")
            choice = input("请输入选项 (1/2): ").strip()
            if choice == '1':
                args.model = 'deepseek'
            elif choice == '2':
                args.model = 'ollama'
            else:
                print("无效的选项，使用默认模型: ollama")
                args.model = 'ollama'
        else:
            # 如果未指定，默认ollama
            if args.model is None:
                args.model = 'ollama'
        
        # 创建NL2SQL实例
        n2s = NL2SQL(model=args.model)
        output_dir = "query_results"
        
        # 示例问题列表
        questions = [
            "查询所有租户的基本信息",
            "显示每个租户的用户数量",
            "查找最近一个月新增的租户",
            "统计每个租户的合同总数",
        ]
        
        # 处理每个问题
        for i, question in enumerate(questions, 1):
            print(f"\n处理第 {i}/{len(questions)} 个问题...")
            try:
                print(f"\n问题: {question}")
                result = n2s.nl2result(question)
                print_result(question, result, output_dir)
            except Exception as e:
                error_result = {
                    'success': False,
                    'error': str(e),
                    'traceback': traceback.format_exc()
                }
                print_result(question, error_result, output_dir)
                
    except Exception as e:
        print(f"程序执行出错: {str(e)}")
        print("\n错误详情:")
        print(traceback.format_exc())

def main_sql_only():
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='自然语言转SQL（仅生成SQL语句）工具')
    parser.add_argument('--model', type=str, choices=['deepseek', 'ollama'],
                        help='选择使用的模型服务: deepseek 或 ollama (默认)')
    parser.add_argument('--interactive', action='store_true', 
                        help='交互式选择模型服务')
    args = parser.parse_args()
    
    # 如果指定了交互式模式或未指定使用的模型服务，让用户选择模型
    if args.interactive or args.model is None:
        print("请选择要使用的模型服务:")
        print("1. DeepSeek API")
        print("2. 本地OLLAMA服务")
        choice = input("请输入选项 (1/2): ").strip()
        if choice == '1':
            args.model = 'deepseek'
        elif choice == '2':
            args.model = 'ollama'
        else:
            print("无效的选项，使用默认模型: ollama")
            args.model = 'ollama'
    else:
        # 如果未指定，默认ollama
        if args.model is None:
            args.model = 'ollama'
    
    # 创建NL2SQL实例
    n2s = NL2SQL(model=args.model)
    
    # 示例问题列表
    questions = [
        "查询所有租户的基本信息",
        # "显示每个租户的用户数量",
        # "查找最近一个月新增的租户",
        "统计每个租户的合同总数",
    ]
    
    # 处理每个问题
    output_dir = "query_results"
    for i, question in enumerate(questions, 1):
        print(f"\n处理第 {i}/{len(questions)} 个问题...")
        try:
            print(f"\n问题: {question}")
            sql = n2s.nl2sql(question)
            print("生成的SQL:")
            print(sql)
            filepath = save_sql_to_file(question, sql, output_dir)
            print(f"\nSQL已保存到文件: {filepath}")
        except Exception as e:
            print(f"生成SQL失败: {str(e)}")
            print(traceback.format_exc())

if __name__ == "__main__":
    main_sql_only() 